from fastapi import FastAPI, Request, Depends
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import time
import os
from pathlib import Path
import asyncio
from typing import Callable, Dict, Any
import logging
import jieba

# 导入环境配置
from kg_retrieval.env_config import APP_CONFIG, MODEL_CONFIG, ES_CONFIG, REDIS_CONFIG, LOGS_DIR, CACHE_DIR

# 设置日志级别
logging.basicConfig(level=getattr(logging, APP_CONFIG["LOG_LEVEL"], logging.INFO))
logger = logging.getLogger("kg-retrieve")

from config import *
from fastapi import APIRouter
from kg_retrieval.monitoring import performance_monitor
from kg_retrieval.es_client import get_es_client, close_es_client
from kg_retrieval.optimized_embedding import OptimizedEmbedding
from kg_retrieval.optimized_rerank import OptimizedReranker

# 创建应用
app = FastAPI(
    title=app_name,
    description="知识库检索系统，提供高性能的文档检索服务",
    version="1.0.0"
)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有源，生产环境建议指定具体域名
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有HTTP方法
    allow_headers=["*"],  # 允许所有头信息
)

# 添加中间件记录请求时长
@app.middleware("http")
async def add_process_time_header(request: Request, call_next: Callable):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response

router_main = APIRouter()

# 定义根路径路由
@router_main.get("/")
def root():
    return {"message": f"欢迎访问{app_name}，API文档请访问 /docs"}

# 健康检查接口
@router_main.get("/health")
async def health_check():
    # 导入router_kg以获取模型注册表
    from router_kg import model_registry
    
    # 检查ES连接
    es_client = await get_es_client()
    es_status = "connected" if es_client else "disconnected"
    
    # 检查模型状态
    embedding_status = "loaded" if model_registry.embedding_model and model_registry.embedding_model.emb_model else "not_loaded"
    rerank_status = "loaded" if model_registry.rerank_model and model_registry.rerank_model.reranker else "not_loaded"
    
    return {
        "status": "healthy",
        "elasticsearch": es_status,
        "embedding_model": embedding_status,
        "rerank_model": rerank_status,
        "version": "1.0.0"
    }

# 注册路由
app.include_router(router_main)

# 设置jieba缓存目录
def setup_jieba_cache():
    """设置jieba缓存目录"""
    # 使用绝对路径
    cache_dir = os.path.abspath(MODEL_CONFIG.get("JIEBA_CACHE_DIR", "./cache"))
    
    # 确保缓存目录存在
    os.makedirs(cache_dir, exist_ok=True)
    
    # 设置jieba缓存路径
    jieba.dt.cache_file = os.path.join(cache_dir, "jieba.cache")
    print(f"设置jieba缓存路径: {jieba.dt.cache_file}")
    
    # 设置jieba临时目录
    jieba.tmp_dir = cache_dir
    print(f"设置jieba临时目录: {jieba.tmp_dir}")
    
    # 预加载jieba
    print("预加载jieba分词...")
    jieba.lcut("预热jieba分词")
    print("jieba分词预加载完成")

# 事件处理
@app.on_event("startup")
async def startup_event():
    """应用启动时执行"""
    print(f"应用 {app_name} 启动中...")
    
    # 设置jieba缓存并预加载
    setup_jieba_cache()
    
    # 初始化ES连接
    await get_es_client()
    
    # 预加载嵌入模型
    print("初始化 OptimizedEmbedding")
    embedding_model = OptimizedEmbedding.preload()
    
    # 预加载重排序模型
    print("初始化 OptimizedReranker")
    rerank_model = OptimizedReranker.preload()
    
    # 将预加载的模型添加到路由模块
    # 导入router_kg并注册模型
    from router_kg import set_models, router_kg
    set_models(embedding_model, rerank_model)
    
    # 注册路由
    app.include_router(router_kg, prefix=router_dict["router_kg"])

@app.on_event("shutdown")
async def shutdown_event():
    """应用关闭时执行"""
    print(f"应用 {app_name} 关闭中...")
    # 关闭ES连接
    await close_es_client()

if __name__ == "__main__":
    import uvicorn
    host = os.environ.get("HOST", APP_CONFIG["HOST"])
    port = int(os.environ.get("PORT", APP_CONFIG["PORT"]))
    
    print(f"启动服务器：http://{host}:{port}")
    print(f"API文档：http://{host}:{port}/docs")
    
    uvicorn.run("main:app", host=host, port=port, reload=True)
